/**
 * @file graph_run_001.cpp
 *
 * Copyright (c) Huawei Technologies Co., Ltd. 2012-2018. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */

#include <iostream>
#include <fstream>
#include <algorithm>
#include <string.h>
#include "graph_utils.h"
#include "utils.h"

using namespace std;
using namespace ge;

bool GenGraph(GraphPtr &graph, AscendString graph_name)
{
    auto shape_data = vector<int64_t>({ 1, 1, 28, 28 });
    TensorDesc desc_data(ge::Shape(shape_data), FORMAT_ND, DT_FLOAT16);

    // data op
    auto data = op::Data("data");
    data.update_input_desc_x(desc_data);
    data.update_output_desc_y(desc_data);
    auto add = op::Add("add").set_input_x1(data).set_input_x2(data);

    std::vector<Operator> inputs { data };
    std::vector<Operator> outputs { add };
    graph =  Graph::ConstructFromInputs(inputs, graph_name);
    graph->SetOutputs(outputs);
    return true;
}


int main(int argc, char *argv[])
{
    // system init
    std::map<AscendString, AscendString> config = { { "ge.exec.deviceId", "0" },
        { "ge.graphRunMode", "1" },
        { "ge.exec.precision_mode", "allow_fp32_to_fp16" } };
    Status ret = ge::GEInitialize(config);
    if (ret != SUCCESS) {
        ERROR_LOG("Initialize ge Failed");
        return OPENCANN_FAILED;
    }
    INFO_LOG("Initialize ge success");

    // Genetate graph
    uint32_t graph_id = 0;
    GraphPtr graph;
    AscendString graph_name("IrGraph");
    ret = GenGraph(graph, graph_name);
    if (!ret) {
        ERROR_LOG("Generate Graph and set input output Failed");
        return OPENCANN_FAILED;
    } else {
        INFO_LOG("Generate Graph and set input output Success");
    }

    // 2. creat session
    std::map<ge::AscendString, ge::AscendString> options;
    ge::Session *session = new Session(options);
    if (session == nullptr) {
        ERROR_LOG("Create session failed");
        return OPENCANN_FAILED;
    } else {
        INFO_LOG("Create session success");
    }

    // 3. add graph
    ret = session->AddGraph(graph_id, *graph);
    if (ret != SUCCESS) {
        ERROR_LOG("add graph failed");
        return OPENCANN_FAILED;
    } else {
        INFO_LOG("add graph success");
    }

    // 4. build graph
    std::vector<Tensor> input_data;
    std::vector<Tensor> output_data;
    vector<int64_t> input_shapes = { 1, 1, 28, 28 };

    const char *workPath = getenv("OPENCANN_TPATH");
    if (workPath == NULL) {
        ERROR_LOG("OPENCANN_TPATH is not set");
        return FAILED;
    }
    std::string input_path = workPath;
    input_path += "/common/test_data/data/ge_add_input_x.bin";
    Tensor input_tensor;
    bool result = GetTensorFromBin(input_path, input_shapes, input_tensor);
    if (!result) {
        ERROR_LOG("Get tensor from bin failed");
        return OPENCANN_FAILED;
    }
    input_data.push_back(input_tensor);
    ret = session->BuildGraph(graph_id, input_data);
    if (ret != SUCCESS) {
        ERROR_LOG("build graph failed");
        return OPENCANN_FAILED;
    } else {
        INFO_LOG("build graph success");
    }
    // 5. Run graph
    ret = session->RunGraph(graph_id, input_data, output_data);
    if (ret != SUCCESS) {
        ERROR_LOG("run graph failed");
        return OPENCANN_FAILED;
    } else {
        INFO_LOG("run graph success");
    }

    // 6. system finalize
    ret = ge::GEFinalize();
    if (ret != SUCCESS) {
        ERROR_LOG("ge finalize failed");
        return OPENCANN_FAILED;
    }
    INFO_LOG("Finalize ge success");

    return OPENCANN_SUCCESS;
}
